import torch
import torch.nn as nn

class DiscriminatorLSTM_no(nn.Module):
    def __init__(self, input_dim, hidden_dim, out_dim):
        super(DiscriminatorLSTM_no, self).__init__()
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers=2)
        self.out = nn.Linear(hidden_dim, out_dim)

    def forward(self, input, lyrics):
        concat_input = torch.cat((input, lyrics), 2)
        # print(concat_input)
        # print(concat_input.shape)
        _, lstm_out = self.lstm(concat_input.view(concat_input.shape[1], concat_input.shape[0], -1))
        ct = lstm_out[1]
        # print(ct.shape)
        # print(ct[-1])
        last_layer_out = ct[-1]
        last_layer_out = last_layer_out.view(last_layer_out.shape[0], 1, -1)

        out = self.out(last_layer_out)
        out = out.view(out.shape[0], -1)
        return out
